#!/usr/bin/python3
"""
OpenCloudOS 9 AI软件自动化验证工具
用于批量验证Python AI软件包在OpenCloudOS 9上的兼容性
"""

import json
import logging
import subprocess
import sys
import time
import traceback
import importlib
# 使用 importlib.metadata 替代 pkg_resources (Python 3.8+)
try:
    from importlib.metadata import distribution, PackageNotFoundError
except ImportError:
    # 兼容 Python < 3.8 的版本
    from importlib_metadata import distribution, PackageNotFoundError
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any
import tempfile
import os
import psutil  # 添加系统资源监控


class AIPackageValidator:
    """AI软件包验证器"""
    
    def __init__(self, log_dir: str = "validation_logs", default_install_timeout: int = 120):
        """
        初始化验证器
        
        Args:
            log_dir: 日志输出目录
            default_install_timeout: 默认安装超时时间（秒）
        """
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(exist_ok=True)
        self.default_install_timeout = default_install_timeout
        
        # 设置日志
        self._setup_logging()
        
        # 验证结果存储
        self.validation_results = {}
        
    def _setup_logging(self):
        """设置日志配置"""
        log_file = self.log_dir / f"validator_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler(sys.stdout)
            ]
        )
        self.logger = logging.getLogger(__name__)
        
    def _monitor_resources(self) -> Dict[str, Any]:
        """监控系统资源使用情况"""
        try:
            cpu_percent = psutil.cpu_percent(interval=1)
            memory = psutil.virtual_memory()
            disk = psutil.disk_usage('/')
            
            return {
                "cpu_percent": cpu_percent,
                "memory_percent": memory.percent,
                "memory_available_gb": round(memory.available / (1024**3), 2),
                "disk_percent": disk.percent,
                "disk_free_gb": round(disk.free / (1024**3), 2)
            }
        except Exception as e:
            self.logger.warning(f"资源监控失败: {e}")
            return {}
    
    def install_package(self, package_name: str, version: str = None, timeout: int = None) -> Dict[str, Any]:
        """
        安装Python包
        
        Args:
            package_name: 包名
            version: 指定版本，如果为None则安装最新版本
            timeout: 安装超时时间（秒），如果为None则使用默认值
            
        Returns:
            安装结果字典
        """
        # 使用传入的超时时间或默认超时时间
        install_timeout = timeout if timeout is not None else self.default_install_timeout
        
        install_result = {
            "package": package_name,
            "version": version,
            "install_success": False,
            "install_time": None,
            "install_log": "",
            "error_message": "",
            "timeout_occurred": False,
            "system_resources_before": self._monitor_resources(),
            "system_resources_after": {}
        }
        
        try:
            # 构建安装命令
            if version:
                package_spec = f"{package_name}=={version}"
            else:
                package_spec = package_name
                
            self.logger.info(f"开始安装包: {package_spec} (超时时间: {install_timeout}秒)")
            
            start_time = time.time()
            
            # 执行pip install
            result = subprocess.run(
                [sys.executable, "-m", "pip", "install", package_spec, "--no-cache-dir"],
                capture_output=True,
                text=True,
                timeout=install_timeout  # 使用指定的超时时间
            )
            
            install_time = time.time() - start_time
            install_result["install_time"] = install_time
            install_result["install_log"] = result.stdout + result.stderr
            install_result["system_resources_after"] = self._monitor_resources()
            
            if result.returncode == 0:
                install_result["install_success"] = True
                self.logger.info(f"成功安装 {package_spec}, 耗时: {install_time:.2f}秒")
            else:
                install_result["error_message"] = f"安装失败，返回码: {result.returncode}"
                self.logger.error(f"安装失败 {package_spec}: {result.stderr}")
                
        except subprocess.TimeoutExpired:
            install_result["error_message"] = f"安装超时（{install_timeout}秒）"
            install_result["timeout_occurred"] = True
            install_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"安装超时: {package_spec}")
        except Exception as e:
            install_result["error_message"] = str(e)
            install_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"安装异常 {package_spec}: {e}")
            
        return install_result
    
    def upgrade_package(self, package_name: str, timeout: int = None) -> Dict[str, Any]:
        """
        升级Python包到最新版本
        
        Args:
            package_name: 包名
            timeout: 升级超时时间（秒），如果为None则使用默认值
            
        Returns:
            升级结果字典
        """
        # 使用传入的超时时间或默认超时时间
        upgrade_timeout = timeout if timeout is not None else self.default_install_timeout
        
        upgrade_result = {
            "package": package_name,
            "upgrade_success": False,
            "upgrade_time": None,
            "upgrade_log": "",
            "error_message": "",
            "timeout_occurred": False,
            "system_resources_before": self._monitor_resources(),
            "system_resources_after": {}
        }
        
        try:
            self.logger.info(f"开始升级包: {package_name} (超时时间: {upgrade_timeout}秒)")
            
            start_time = time.time()
            
            # 执行pip upgrade
            result = subprocess.run(
                [sys.executable, "-m", "pip", "install", "--upgrade", package_name],
                capture_output=True,
                text=True,
                timeout=upgrade_timeout
            )
            
            upgrade_time = time.time() - start_time
            upgrade_result["upgrade_time"] = upgrade_time
            upgrade_result["upgrade_log"] = result.stdout + result.stderr
            upgrade_result["system_resources_after"] = self._monitor_resources()
            
            if result.returncode == 0:
                upgrade_result["upgrade_success"] = True
                self.logger.info(f"成功升级 {package_name}, 耗时: {upgrade_time:.2f}秒")
            else:
                upgrade_result["error_message"] = f"升级失败，返回码: {result.returncode}"
                self.logger.error(f"升级失败 {package_name}: {result.stderr}")
                
        except subprocess.TimeoutExpired:
            upgrade_result["error_message"] = f"升级超时（{upgrade_timeout}秒）"
            upgrade_result["timeout_occurred"] = True
            upgrade_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"升级超时: {package_name}")
        except Exception as e:
            upgrade_result["error_message"] = str(e)
            upgrade_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"升级异常 {package_name}: {e}")
            
        return upgrade_result
    
    def get_package_info(self, package_name: str) -> Dict[str, Any]:
        """
        获取已安装包的信息
        
        Args:
            package_name: 包名
            
        Returns:
            包信息字典
        """
        try:
            dist = distribution(package_name)
            return {
                "package": package_name,
                "version": dist.version,
                "location": str(getattr(dist, 'location', 'unknown')),
                "requires": [str(req) for req in getattr(dist, 'requires', []) or []]
            }
        except PackageNotFoundError:
            return {"package": package_name, "error": "包未找到"}
        except Exception as e:
            return {"package": package_name, "error": str(e)}
    
    def run_basic_import_test(self, package_name: str, timeout: int = 60) -> Dict[str, Any]:
        """
        执行基本导入测试
        
        Args:
            package_name: 包名
            timeout: 测试超时时间（秒）
            
        Returns:
            测试结果字典
        """
        test_result = {
            "test_type": "basic_import",
            "package": package_name,
            "success": False,
            "execution_time": None,
            "error_message": "",
            "traceback": "",
            "timeout_occurred": False,
            "system_resources_before": self._monitor_resources(),
            "system_resources_after": {}
        }
        
        try:
            start_time = time.time()
            
            # 创建临时测试文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
                f.write(f"import {package_name}\nprint(f'成功导入 {package_name}')")
                temp_file = f.name
            
            try:
                # 执行导入测试
                result = subprocess.run(
                    [sys.executable, temp_file],
                    capture_output=True,
                    text=True,
                    timeout=timeout
                )
                
                execution_time = time.time() - start_time
                test_result["execution_time"] = execution_time
                test_result["system_resources_after"] = self._monitor_resources()
                
                if result.returncode == 0:
                    test_result["success"] = True
                    self.logger.info(f"基本导入测试通过: {package_name}")
                else:
                    test_result["error_message"] = result.stderr
                    self.logger.error(f"导入测试失败 {package_name}: {result.stderr}")
                    
            finally:
                # 清理临时文件
                os.unlink(temp_file)
                
        except subprocess.TimeoutExpired:
            test_result["error_message"] = f"导入测试超时（{timeout}秒）"
            test_result["timeout_occurred"] = True
            test_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"导入测试超时: {package_name}")
        except Exception as e:
            test_result["error_message"] = str(e)
            test_result["traceback"] = traceback.format_exc()
            test_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"导入测试异常 {package_name}: {e}")
            
        return test_result
    
    def run_custom_test(self, package_name: str, test_case: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行自定义测试用例
        
        Args:
            package_name: 包名
            test_case: 测试用例配置
            
        Returns:
            测试结果字典
        """
        test_result = {
            "test_type": "custom",
            "package": package_name,
            "test_name": test_case.get("name", "unnamed_test"),
            "success": False,
            "execution_time": None,
            "error_message": "",
            "traceback": "",
            "output": "",
            "timeout_occurred": False,
            "system_resources_before": self._monitor_resources(),
            "system_resources_after": {}
        }
        
        try:
            start_time = time.time()
            
            # 创建临时测试文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
                f.write(test_case.get("code", ""))
                temp_file = f.name
            
            try:
                # 执行测试代码
                result = subprocess.run(
                    [sys.executable, temp_file],
                    capture_output=True,
                    text=True,
                    timeout=test_case.get("timeout", 60)
                )
                
                execution_time = time.time() - start_time
                test_result["execution_time"] = execution_time
                test_result["output"] = result.stdout
                test_result["system_resources_after"] = self._monitor_resources()
                
                if result.returncode == 0:
                    test_result["success"] = True
                    self.logger.info(f"自定义测试通过: {package_name} - {test_case.get('name')}")
                else:
                    test_result["error_message"] = result.stderr
                    self.logger.error(f"自定义测试失败 {package_name}: {result.stderr}")
                    
            finally:
                # 清理临时文件
                os.unlink(temp_file)
                
        except subprocess.TimeoutExpired:
            test_result["error_message"] = f"测试超时（{test_case.get('timeout', 60)}秒）"
            test_result["timeout_occurred"] = True
            test_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"自定义测试超时: {package_name}")
        except Exception as e:
            test_result["error_message"] = str(e)
            test_result["traceback"] = traceback.format_exc()
            test_result["system_resources_after"] = self._monitor_resources()
            self.logger.error(f"自定义测试异常 {package_name}: {e}")
            
        return test_result
    
    def validate_package(self, package_config: Dict[str, Any], upgrade_before_test: bool = False) -> Dict[str, Any]:
        """
        验证单个包
        
        Args:
            package_config: 包配置字典
            upgrade_before_test: 是否在测试前升级包
            
        Returns:
            完整的验证结果
        """
        package_name = package_config.get("name")
        version = package_config.get("version")
        
        self.logger.info(f"开始验证包: {package_name}")
        
        validation_result = {
            "package": package_name,
            "version": version,
            "timestamp": datetime.now().isoformat(),
            "opencloudos_version": "9",
            "install_result": {},
            "upgrade_result": {},
            "package_info": {},
            "tests": [],
            "overall_success": False,
            "timeout_occurred": False,
            "system_resources_initial": self._monitor_resources()
        }
        
        # 1. 如果需要，先升级包
        if upgrade_before_test:
            upgrade_timeout = package_config.get("upgrade_timeout", self.default_install_timeout)
            upgrade_result = self.upgrade_package(package_name, timeout=upgrade_timeout)
            validation_result["upgrade_result"] = upgrade_result
            
            # 如果升级超时，标记超时并继续
            if upgrade_result.get("timeout_occurred", False):
                validation_result["timeout_occurred"] = True
                validation_result["system_resources_final"] = self._monitor_resources()
                self._save_package_log(validation_result)
                return validation_result
        
        # 2. 安装包（如果指定了版本）
        install_timeout = package_config.get("install_timeout", self.default_install_timeout)
        install_result = self.install_package(package_name, version, install_timeout)
        validation_result["install_result"] = install_result
        
        # 如果安装超时，标记超时并继续
        if install_result.get("timeout_occurred", False):
            validation_result["timeout_occurred"] = True
            validation_result["system_resources_final"] = self._monitor_resources()
            self._save_package_log(validation_result)
            return validation_result
            
        if not install_result["install_success"]:
            validation_result["overall_success"] = False
            validation_result["system_resources_final"] = self._monitor_resources()
            self._save_package_log(validation_result)
            return validation_result
        
        # 3. 获取包信息
        validation_result["package_info"] = self.get_package_info(package_name)
        
        # 4. 执行基本导入测试
        import_test_timeout = package_config.get("import_test_timeout", 60)
        import_test = self.run_basic_import_test(package_name, timeout=import_test_timeout)
        validation_result["tests"].append(import_test)
        
        # 如果导入测试超时，标记超时并继续
        if import_test.get("timeout_occurred", False):
            validation_result["timeout_occurred"] = True
            validation_result["system_resources_final"] = self._monitor_resources()
            self._save_package_log(validation_result)
            return validation_result
        
        # 5. 执行自定义测试
        custom_tests = package_config.get("custom_tests", [])
        for test_case in custom_tests:
            custom_test_result = self.run_custom_test(package_name, test_case)
            validation_result["tests"].append(custom_test_result)
            
            # 如果任何自定义测试超时，标记超时并继续
            if custom_test_result.get("timeout_occurred", False):
                validation_result["timeout_occurred"] = True
                validation_result["system_resources_final"] = self._monitor_resources()
                self._save_package_log(validation_result)
                return validation_result
        
        # 6. 判断整体成功状态
        all_tests_passed = all(test["success"] for test in validation_result["tests"])
        validation_result["overall_success"] = install_result["install_success"] and all_tests_passed
        
        # 7. 保存单个包的日志
        validation_result["system_resources_final"] = self._monitor_resources()
        self._save_package_log(validation_result)
        
        return validation_result
    
    def _save_package_log(self, validation_result: Dict[str, Any]):
        """保存单个包的验证日志"""
        package_name = validation_result["package"]
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = self.log_dir / f"{package_name}_{timestamp}.json"
        
        with open(log_file, 'w', encoding='utf-8') as f:
            json.dump(validation_result, f, ensure_ascii=False, indent=2)
        
        self.logger.info(f"已保存验证日志: {log_file}")
    
    def batch_validate(self, packages_config: List[Dict[str, Any]], parallel: bool = False, upgrade_before_test: bool = False) -> Dict[str, Any]:
        """
        批量验证包
        
        Args:
            packages_config: 包配置列表
            parallel: 是否并行执行（当前版本为串行，后续可扩展）
            upgrade_before_test: 是否在测试前升级所有包
            
        Returns:
            批量验证结果汇总
        """
        batch_result = {
            "batch_timestamp": datetime.now().isoformat(),
            "total_packages": len(packages_config),
            "successful_packages": 0,
            "failed_packages": 0,
            "timeout_packages": 0,
            "packages_results": [],
            "summary": {},
            "system_resources_initial": self._monitor_resources()
        }
        
        self.logger.info(f"开始批量验证 {len(packages_config)} 个包")
        
        for package_config in packages_config:
            try:
                result = self.validate_package(package_config, upgrade_before_test)
                batch_result["packages_results"].append(result)
                
                if result.get("timeout_occurred", False):
                    batch_result["timeout_packages"] += 1
                elif result["overall_success"]:
                    batch_result["successful_packages"] += 1
                else:
                    batch_result["failed_packages"] += 1
                    
            except Exception as e:
                error_result = {
                    "package": package_config.get("name", "unknown"),
                    "error": str(e),
                    "traceback": traceback.format_exc(),
                    "overall_success": False
                }
                batch_result["packages_results"].append(error_result)
                batch_result["failed_packages"] += 1
                self.logger.error(f"验证包时发生异常: {e}")
        
        # 生成汇总信息
        batch_result["summary"] = {
            "success_rate": batch_result["successful_packages"] / batch_result["total_packages"] * 100,
            "timeout_packages_list": [
                result["package"] for result in batch_result["packages_results"] 
                if result.get("timeout_occurred", False)
            ],
            "failed_packages_list": [
                result["package"] for result in batch_result["packages_results"] 
                if not result.get("overall_success", False) and not result.get("timeout_occurred", False)
            ]
        }
        
        # 保存批量验证结果
        batch_result["system_resources_final"] = self._monitor_resources()
        self._save_batch_log(batch_result)
        
        self.logger.info(f"批量验证完成: 成功 {batch_result['successful_packages']}/{batch_result['total_packages']}")
        self.logger.info(f"超时包: {batch_result['timeout_packages']}")
        self.logger.info(f"失败包: {batch_result['failed_packages']}")
        
        return batch_result
    
    def _save_batch_log(self, batch_result: Dict[str, Any]):
        """保存批量验证日志"""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = self.log_dir / f"batch_validation_{timestamp}.json"
        
        with open(log_file, 'w', encoding='utf-8') as f:
            json.dump(batch_result, f, ensure_ascii=False, indent=2)
        
        self.logger.info(f"已保存批量验证日志: {log_file}")


def main():
    """主函数示例"""
    validator = AIPackageValidator()
    
    # 示例配置
    example_packages = [
        {
            "name": "numpy",
            "version": None,  # 最新版本
            "custom_tests": [
                {
                    "name": "basic_numpy_test",
                    "code": """
import numpy as np
# 基本功能测试
arr = np.array([1, 2, 3, 4, 5])
print(f"Array: {arr}")
print(f"Mean: {np.mean(arr)}")
print(f"Shape: {arr.shape}")
assert len(arr) == 5
print("NumPy基本测试通过!")
""",
                    "timeout": 30
                }
            ]
        },
        {
            "name": "pandas",
            "version": None,
            "custom_tests": [
                {
                    "name": "basic_pandas_test",
                    "code": """
import pandas as pd
import numpy as np
# 基本功能测试
df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
print(f"DataFrame:\\n{df}")
print(f"Shape: {df.shape}")
assert df.shape == (3, 2)
print("Pandas基本测试通过!")
""",
                    "timeout": 30
                }
            ]
        }
    ]
    
    # 执行批量验证
    result = validator.batch_validate(example_packages)
    
    print("\n=== 验证结果汇总 ===")
    print(f"总包数: {result['total_packages']}")
    print(f"成功: {result['successful_packages']}")
    print(f"超时: {result['timeout_packages']}")
    print(f"失败: {result['failed_packages']}")
    print(f"成功率: {result['summary']['success_rate']:.1f}%")
    
    if result['summary']['timeout_packages_list']:
        print(f"超时的包: {', '.join(result['summary']['timeout_packages_list'])}")
    
    if result['summary']['failed_packages_list']:
        print(f"失败的包: {', '.join(result['summary']['failed_packages_list'])}")


if __name__ == "__main__":
    main()